# -*-coding: utf-8 -*-
"""
    @Theme   : Setup and check confirm merge
    @Time    : 2023/12/14 21:24
    @Author  : YamFish
    @Url     : https://www.zhihu.com/people/yamfish
"""
import os
import re
import json
import time
import random
import string
import collections
from threading import Thread
from tempfile import TemporaryDirectory

import click


def write_random_data(path, write_count, interval):
    with open(path, 'wt') as f:
        for i in range(write_count):
            time.sleep(random.random() * interval)
            letters = random.choices(string.ascii_lowercase, k=10)
            data = f'{path}-{i:02}-{"".join(letters)}\n'
            f.write(data)
            f.flush()


def start_write_threads(directory, file_count):
    paths = []
    for i in range(file_count):
        path = os.path.join(directory, str(i))
        with open(path, 'w'):
            # Make sure the file at this path will exist when
            # the reading thread tries to poll it.
            pass
        paths.append(path)
        thread = Thread(target=write_random_data, args=(path, 10, 0.1))
        thread.start()
    return paths


def close_all(handles):
    # Make Sure the write threads have finished
    time.sleep(1.5)
    for handle in handles:
        handle.close()


def setup(count):
    tmp_dir = TemporaryDirectory()

    write_paths = start_write_threads(tmp_dir.name, count)
    merged_path = os.path.join(tmp_dir.name, 'merged')

    file_handles = [open(path, 'rb') for path in write_paths]
    Thread(target=close_all, args=(file_handles,)).start()

    return tmp_dir, write_paths, file_handles, merged_path


def confirm_merge(write_paths, merged_path, verbose=False):
    merged = collections.defaultdict(list)
    with open(merged_path, 'rt') as f:
        for line in f:
            path = re.match(r'(.*?)-\d{2}-', line)
            if path and path.group(1) in write_paths:
                merged[path.group(1)].append(line)

    expected = collections.defaultdict(list)
    for path in write_paths:
        with open(path, 'rt') as f:
            expected[path].extend(f.readlines())

    for key, expected_lines in expected.items():
        found_lines = merged[key]
        assert expected_lines == found_lines

    if verbose:
        print('Merged: ', json.dumps(merged, indent=2))
        print('Expected: ', json.dumps(expected, indent=2))


@click.command()
@click.option('--sleep', default=60, help='you kan check the files in tmpdir.', type=float)
@click.option('--files', default=5, help='the number of files.', type=int)
@click.option('--verbose', is_flag=True, help='show more details.')
def main(sleep, verbose, files):
    """The specified number of files will be generated in a temporary directory."""
    tmpdir, input_paths, handles, output_path = setup(files)
    print(f'tmpdir is {tmpdir.name}')

    if verbose:
        print(f'writing paths: \n'
              f'{os.linesep.join(input_paths)}')
        print(f'the merged path should be: {output_path} ')

    time.sleep(sleep)
    tmpdir.cleanup()


if __name__ == "__main__":
    main()
